import argparse

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--eval_batch_size', type=int, default=64)
    parser.add_argument('--val_steps', type=int, default=10)
    parser.add_argument('--cuda', type=str, default='0')
    parser.add_argument('--dataname', type=str, default='mimic')
    parser.add_argument('--hdim', type=int, default=64)
    parser.add_argument('--title', type=str, default='test')
    parser.add_argument('--start_epoch', type=int, default=0)
    parser.add_argument('--max_epoch', type=int, default=200)
    parser.add_argument('--load_path', type=str, default=None)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--warmups', type=int, default=0)
    parser.add_argument('--method', type=str, default='attnhp')
    parser.add_argument('--components', type=int, default=64)
    parser.add_argument('--horizon', type=str, default='5.,5.,5.')
    parser.add_argument('--omega', type=float, default=1.)
    parser.add_argument('--layers', type=int, default=1)
    parser.add_argument('--nhead', type=int, default=2)
    parser.add_argument('--beta', type=float, default=1.)
    parser.add_argument('--siren_layers', type=int, default=1)
    # parser.add_argument('--dec_type', type=str, default='lnm')
    parser.add_argument('--max_len', type=int, default=200)
    # parser.add_argument('--attn_type', type=str, default='regular')
    parser.add_argument('--reg', type=float, default=0)
    parser.add_argument('--mlr', type=float, default=1e-3)
    parser.add_argument('--weight_decay', type=float, default=1e-4)
    parser.add_argument('--max_train_seqs', type=int, default=-1)
    parser.add_argument('--max_dt', type=float, default=-1)
    parser.add_argument('--optim', type=str, default='adamw')
    parser.add_argument('--num_steps', type=int, default=200)
    args = parser.parse_args()
    assert args.warmups < args.max_epoch
    
    return args

def load_parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--title', type=str, default='hawkes-nhp')
    parser.add_argument('--test_batch_size', type=int, default=64)
    parser.add_argument('--cuda', type=str, default='1')
    parser.add_argument('--hist_len', type=int, default=20)
    parser.add_argument('--pred_len', type=int, default=10)
    parser.add_argument('--size', type=int, default=-1)
    parser.add_argument('--stride', type=int, default=10)
    args = parser.parse_args()
    return args